import streamlit as st
import os
from docx import Document
from docx.shared import Inches, Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from PyPDF2 import PdfReader
import tempfile
import requests
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import math
from io import BytesIO
import re

class Translator:
    def __init__(self):
        self.api_url = 'https://api.myhispreadnlp.com/v1/chat/completions'
        self.model = "gpt-4o"
        self.max_retries = 3
        self.retry_delay = 1  # 减少重试延迟
    
    def translate_with_retry(self, api_key, text, target_language):
        """带重试机制的翻译"""
        for attempt in range(self.max_retries):
            try:
                return self.translate(api_key, text, target_language)
            except Exception as e:
                if attempt == self.max_retries - 1:
                    raise e
                time.sleep(self.retry_delay * (attempt + 1))
    
    def translate(self, api_key, text, target_language):
        """翻译文本到目标语言"""
        headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }

        data = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": "You are a professional translator."},
                {"role": "user", "content": f"Translate this text to {target_language}: {text}"}
            ]
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            return result['choices'][0]['message']['content']
        except requests.exceptions.RequestException as e:
            raise Exception(f"API请求错误: {str(e)}")
        except (KeyError, IndexError) as e:
            raise Exception(f"API响应格式错误: {str(e)}")

class TextProcessor:
    def __init__(self):
        self.max_chunk_size = 6800  # 每块最大字符数
        self.overlap_size = 200  # 增加重叠区域，确保上下文连贯
        self.min_chunk_size = 1000  # 最小块大小
    
    def split_text_into_chunks(self, text):
        """将文本分割成适当大小的块，带有重叠区域"""
        chunks = []
        start = 0
        text_length = len(text)

        while start < text_length:
            # 计算当前块的结束位置
            end = min(start + self.max_chunk_size, text_length)
            
            if end >= text_length:
                chunk = text[start:]
                if len(chunk) >= self.min_chunk_size:
                    chunks.append(chunk)
                elif chunks:  # 如果最后一块太小，附加到前一块
                    chunks[-1] = chunks[-1] + chunk
                else:  # 如果这是唯一的块，直接添加
                    chunks.append(chunk)
                break
            
            # 在句子边界处分割
            split_position = self._find_sentence_boundary(text[start:end])
            
            if split_position > 0:
                chunk = text[start:start + split_position]
                if len(chunk) >= self.min_chunk_size:
                    chunks.append(chunk)
                    start = start + split_position - self.overlap_size
                else:  # 如果块太小，扩展到下一个句子边界
                    next_boundary = self._find_next_sentence_boundary(text[start + split_position:])
                    if next_boundary > 0:
                        chunks.append(text[start:start + split_position + next_boundary])
                        start = start + split_position + next_boundary - self.overlap_size
                    else:
                        chunks.append(text[start:end])
                        start = end - self.overlap_size
            else:
                chunks.append(text[start:end])
                start = end - self.overlap_size
        
        return chunks

    def _find_sentence_boundary(self, text):
        """在文本中找到最后一个完整句子的位置"""
        sentence_endings = ['. ', '。', '！', '? ', '？', '\n', ';', '；']
        last_position = 0
        
        for ending in sentence_endings:
            pos = text.rfind(ending)
            if pos > last_position:
                last_position = pos + len(ending)
        
        return last_position if last_position > 0 else len(text)

    def _find_next_sentence_boundary(self, text):
        """找到下一个句子边界"""
        sentence_endings = ['. ', '。', '！', '? ', '？', '\n', ';', '；']
        first_position = len(text)
        
        for ending in sentence_endings:
            pos = text.find(ending)
            if pos != -1 and pos < first_position:
                first_position = pos + len(ending)
        
        return first_position if first_position < len(text) else 0

def validate_translation(text, translated_text):
    """验证翻译结果是否有效"""
    if not translated_text or len(translated_text.strip()) == 0:
        return False, "空翻译结果"
    
    source_length = len(text.strip())
    translated_length = len(translated_text.strip())
    
    # 更宽松的长度比例验证
    if source_length < 100:
        min_ratio = 0.2  # 降低最小比例
        max_ratio = 4.0  # 提高最大比例
    elif source_length < 500:
        min_ratio = 0.15
        max_ratio = 5.0
    else:
        min_ratio = 0.1
        max_ratio = 6.0
    
    # 检查长度比例
    length_ratio = translated_length / source_length
    if length_ratio < min_ratio:
        return False, f"翻译结果过短 (比例: {length_ratio:.2f})"
    if length_ratio > max_ratio:
        return False, f"翻译结果过长 (比例: {length_ratio:.2f})"
    
    # 更智能的标点符号检查
    if len(translated_text) > 100:  # 增加长度阈值
        basic_punctuation = ['.', '。', '!', '！', '?', '？', ',', '，', '、', ';', '；']
        sentence_count = sum(translated_text.count(p) for p in ['.', '。', '!', '！', '?', '？'])
        if sentence_count == 0:
            return False, "缺少句子结束标点"
    
    return True, "有效"

def translate_chunk(args):
    """翻译单个文本块的函数"""
    translator, api_key, chunk, target_language, chunk_index, total_chunks = args
    max_attempts = 5
    base_delay = 1
    
    if not chunk or len(chunk.strip()) == 0:
        return chunk_index, None, "空文本块"
    
    original_chunk = chunk
    for attempt in range(max_attempts):
        try:
            if attempt > 0:
                delay = base_delay * (1.5 ** attempt)
                time.sleep(delay)
            
            # 对较长文本进行预处理
            if len(chunk) > 3000 and attempt > 2:
                # 尝试清理和规范化文本
                chunk = chunk.replace('\n\n', '\n').strip()
            
            translated = translator.translate_with_retry(api_key, chunk, target_language)
            valid, message = validate_translation(chunk, translated)
            
            if valid:
                return chunk_index, translated, None
            else:
                if attempt == max_attempts - 1:
                    # 最后一次尝试，分割处理
                    if len(chunk) > 1000:
                        splits = []
                        # 按句子分割
                        sentences = re.split(r'([.。!！?？]\s*)', chunk)
                        current_part = ""
                        
                        for i in range(0, len(sentences), 2):
                            sentence = sentences[i]
                            ending = sentences[i + 1] if i + 1 < len(sentences) else ""
                            if len(current_part) + len(sentence) + len(ending) < 2000:
                                current_part += sentence + ending
                            else:
                                if current_part:
                                    splits.append(current_part)
                                current_part = sentence + ending
                        
                        if current_part:
                            splits.append(current_part)
                        
                        # 翻译各个部分
                        translated_parts = []
                        for part in splits:
                            try:
                                time.sleep(1)
                                part_translated = translator.translate_with_retry(api_key, part, target_language)
                                valid, _ = validate_translation(part, part_translated)
                                if valid:
                                    translated_parts.append(part_translated)
                            except Exception:
                                continue
                        
                        if translated_parts:
                            combined = " ".join(translated_parts)
                            valid, _ = validate_translation(original_chunk, combined)
                            if valid:
                                return chunk_index, combined, None
                    
                    return chunk_index, None, f"翻译结果无效: {message}"
                
        except Exception as e:
            if attempt == max_attempts - 1:
                return chunk_index, None, str(e)
    
    return chunk_index, None, "达到最大重试次数"

def retry_failed_chunks(translator, api_key, chunks, language, failed_chunks, status_container, progress_bar):
    """重试失败的文本块"""
    retry_results = []
    total_failed = len(failed_chunks)
    max_workers = min(3, total_failed)  # 动态调整并发数
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        retry_tasks = [
            (translator, api_key, chunks[chunk_index], language, chunk_index, total_failed)
            for chunk_index in failed_chunks
        ]
        
        futures = [executor.submit(translate_chunk, args) for args in retry_tasks]
        
        for i, future in enumerate(as_completed(futures)):
            chunk_index, result, error = future.result()
            progress = (i + 1) / total_failed
            progress_bar.progress(min(0.9 + progress * 0.1, 1.0))
            
            if error is None:
                retry_results.append((chunk_index, result))
                with status_container:
                    st.markdown(f"✅ 重试成功：第 {chunk_index + 1} 部分")
            else:
                # 如果重试失败，尝试使用更小的块大小
                try:
                    chunk = chunks[chunk_index]
                    if len(chunk) > 1000:
                        # 分割成更小的块
                        sentences = chunk.split('. ')
                        if len(sentences) > 1:
                            translated_parts = []
                            for sentence in sentences:
                                if sentence.strip():
                                    time.sleep(1)  # 添加延迟避免频率限制
                                    trans = translator.translate_with_retry(api_key, sentence.strip(), language)
                                    if trans:
                                        translated_parts.append(trans)
                            
                            if translated_parts:
                                combined_translation = '. '.join(translated_parts)
                                if validate_translation(chunk, combined_translation)[0]:
                                    retry_results.append((chunk_index, combined_translation))
                                    with status_container:
                                        st.markdown(f"✅ 分句重试成功：第 {chunk_index + 1} 部分")
                                    continue
                except Exception as e:
                    pass
                
                with status_container:
                    st.error(f"❌ 重试失败：第 {chunk_index + 1} 部分 - {error}")
    
    return retry_results

class FileHandler:
    @staticmethod
    def read_docx(file_path):
        """读取Word文档内容"""
        try:
            document = Document(file_path)
            return "\n".join(para.text for para in document.paragraphs)
        except Exception as e:
            raise Exception(f"Word文档读取错误: {str(e)}")

    @staticmethod
    def read_txt(file_path):
        """读取文本文件内容"""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            raise Exception(f"文本文件读取错误: {str(e)}")

    @staticmethod
    def read_pdf(file_path):
        """读取PDF文件内容"""
        try:
            pdf = PdfReader(file_path)
            return "\n".join(page.extract_text() for page in pdf.pages)
        except Exception as e:
            raise Exception(f"PDF文件读取错误: {str(e)}")

def create_word_document(text, style='Normal'):
    """创建Word文档并添加格式化文本"""
    doc = Document()
    
    # 设置页面边距（单位：英寸）
    sections = doc.sections
    for section in sections:
        section.left_margin = Inches(1)
        section.right_margin = Inches(1)
        section.top_margin = Inches(1)
        section.bottom_margin = Inches(1)
    
    # 按段落分割文本并添加到文档
    paragraphs = text.split('\n')
    for para_text in paragraphs:
        if para_text.strip():  # 跳过空段落
            paragraph = doc.add_paragraph()
            paragraph.style = style
            run = paragraph.add_run(para_text)
            # 设置字体
            font = run.font
            font.name = 'Times New Roman'  # 英文字体
            font.size = Pt(12)  # 字号
            paragraph.paragraph_format.space_after = Pt(12)  # 段后间距
            
    return doc

def save_word_document(doc):
    """将文档转换为字节流"""
    buffer = BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer

def main():
    # 设置页面配置
    st.set_page_config(
        page_title="文件翻译助手",
        page_icon="🌐",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    # 使用CSS美化界面
    st.markdown("""
        <style>
        .main {
            padding: 2rem;
        }
        .stButton>button {
            width: 100%;
            height: 3rem;
            margin-top: 1rem;
            margin-bottom: 1rem;
        }
        .upload-text {
            color: #666;
            font-size: 0.9em;
        }
        .success-text {
            color: #28a745;
        }
        .warning-text {
            color: #ffc107;
        }
        </style>
    """, unsafe_allow_html=True)

    # 页面标题
    st.title("🌐 文件翻译助手")
    st.markdown("---")

    # 初始化组件
    translator = Translator()
    file_handler = FileHandler()
    text_processor = TextProcessor()
    
    # API密钥输入
    api_key = st.text_input(
        "🔑 API密钥",
        type="password",
        help="请输入您的API密钥以使用翻译服务"
    )
    
    # 文件上传
    st.markdown("<p class='upload-text'>📁 支持的文件格式: .docx, .txt, .pdf</p>", unsafe_allow_html=True)
    upload = st.file_uploader("选择要翻译的文件", type=["docx", "txt", "pdf"])
    
    # 语言选择
    if upload:
        language = st.selectbox(
            "🌍 选择目标语言",
            ["中文", "英文", "日文", "韩文", "法文", "德文", "俄文", "西班牙文"],
            help="选择您想要翻译成的目标语言"
        )
        
        # 添加开始翻译按钮
        start_translation = st.button("🚀 开始翻译", use_container_width=True)
        
        if start_translation and api_key:
            try:
                # 创建状态容器
                status_container = st.container()
                progress_bar = st.progress(0)
                
                with status_container:
                    st.markdown("#### 📊 处理进度")
                
                # 保存上传文件到临时文件
                with tempfile.NamedTemporaryFile(delete=False, suffix=f".{upload.name.split('.')[-1]}") as tmp:
                    tmp.write(upload.getvalue())
                    file_path = tmp.name

                # 读取文件内容
                with status_container:
                    st.markdown("� 正在读取文件...")
                progress_bar.progress(10)
                
                if upload.name.endswith('.docx'):
                    content = file_handler.read_docx(file_path)
                elif upload.name.endswith('.txt'):
                    content = file_handler.read_txt(file_path)
                elif upload.name.endswith('.pdf'):
                    content = file_handler.read_pdf(file_path)
                
                # 分割文本
                chunks = text_processor.split_text_into_chunks(content)
                total_chunks = len(chunks)
                
                with status_container:
                    st.markdown(f"📑 文件已分割为 {total_chunks} 个部分")
                    st.markdown("🔄 正在翻译...")
                
                # 使用线程池进行并行翻译
                translated_chunks = [""] * total_chunks
                failed_chunks = []
                all_errors = {}  # 存储所有错误信息
                
                with ThreadPoolExecutor(max_workers=5) as executor:
                    translation_tasks = [
                        (translator, api_key, chunk, language, i, total_chunks)
                        for i, chunk in enumerate(chunks)
                    ]
                    
                    futures = [executor.submit(translate_chunk, args) for args in translation_tasks]
                    
                    for i, future in enumerate(as_completed(futures)):
                        chunk_index, result, error = future.result()
                        progress = (i + 1) / total_chunks
                        progress_bar.progress(min(0.1 + progress * 0.8, 0.9))  # 留出重试的进度空间
                        
                        if error is not None:
                            failed_chunks.append(chunk_index)
                            all_errors[chunk_index] = error
                            with status_container:
                                st.error(f"❌ 第 {chunk_index + 1} 部分翻译失败: {error}")
                        else:
                            translated_chunks[chunk_index] = result
                            with status_container:
                                st.markdown(f"✅ 已完成第 {chunk_index + 1}/{total_chunks} 部分")
                
                # 自动重试失败的部分
                if failed_chunks:
                    with status_container:
                        st.markdown("🔄 正在重试失败的部分...")
                    
                    retry_results = retry_failed_chunks(
                        translator, api_key, chunks, language,
                        failed_chunks, status_container, progress_bar
                    )
                    
                    # 更新翻译结果
                    for chunk_index, result in retry_results:
                        if result:
                            translated_chunks[chunk_index] = result
                            failed_chunks.remove(chunk_index)
                
                # 处理最终结果
                if not failed_chunks:
                    # 验证所有翻译结果
                    invalid_chunks = []
                    for i, chunk in enumerate(translated_chunks):
                        if not validate_translation(chunks[i], chunk)[0]:
                            invalid_chunks.append(i)
                    
                    if invalid_chunks:
                        st.error(f"❌ 发现 {len(invalid_chunks)} 个部分翻译结果无效")
                        with st.expander("查看无效翻译详情"):
                            for chunk_index in invalid_chunks:
                                st.error(f"第 {chunk_index + 1} 部分翻译结果无效")
                        failed_chunks.extend(invalid_chunks)
                    else:
                        final_translation = "\n".join(translated_chunks)
                        with status_container:
                            st.markdown("✨ 翻译完成！")
                        
                        st.markdown("### 📝 翻译结果")
                        st.text_area("", final_translation, height=400)
                        
                        # 创建下载按钮列
                        col1, col2 = st.columns(2)
                        
                        with col1:
                            # TXT格式下载
                            st.download_button(
                                label="📥 下载TXT格式",
                                data=final_translation.encode('utf-8'),
                                file_name=f"translated_{upload.name.split('.')[0]}.txt",
                                mime="text/plain",
                                help="下载为文本文件"
                            )
                        
                        with col2:
                            # Word格式下载
                            try:
                                # 创建Word文档
                                doc = create_word_document(final_translation)
                                docx_buffer = save_word_document(doc)
                                
                                st.download_button(
                                    label="📥 下载Word格式",
                                    data=docx_buffer,
                                    file_name=f"translated_{upload.name.split('.')[0]}.docx",
                                    mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                                    help="下载为Word文档"
                                )
                            except Exception as e:
                                st.error(f"Word文档生成失败: {str(e)}")

                else:
                    # 显示失败信息并提供重试选项
                    st.error(f"❌ 仍有 {len(failed_chunks)} 个部分翻译失败")
                    
                    # 显示失败详情
                    with st.expander("查看失败详情"):
                        for chunk_index in failed_chunks:
                            st.error(f"第 {chunk_index + 1} 部分: {all_errors[chunk_index]}")
                    
                    # 提供部分结果下载
                    st.warning("⚠️ 您可以下载已翻译的部分，或者重试失败的部分")
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        partial_translation = "\n".join(
                            chunk for chunk in translated_chunks if chunk
                        )
                        st.download_button(
                            label="📥 下载已翻译部分",
                            data=partial_translation.encode('utf-8'),
                            file_name=f"partial_translated_{upload.name.split('.')[0]}.txt",
                            mime="text/plain",
                        )
                    
                    with col2:
                        if st.button("🔄 重试失败的部分"):
                            st.experimental_rerun()

            except Exception as e:
                st.error(f"❌ 处理过程中出现错误: {str(e)}")
            finally:
                # 清理临时文件
                if 'file_path' in locals():
                    os.remove(file_path)
        elif start_translation and not api_key:
            st.warning("⚠️ 请先输入API密钥")

if __name__ == "__main__":
    main()